import av
import cv2
import time
from fractions import Fraction

# 打开摄像头
cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

# 创建输出容器和视频流
container = av.open("02.mp4", mode="w")
stream = container.add_stream("h264", rate=None)
stream.width = 640
stream.height = 480
stream.pix_fmt = "yuv420p"

# 使用 H.264 通常的时间基：90kHz（1/90000秒）
stream.time_base = Fraction(1, 90000)

# 起始时间
start_time_ms = int(time.time() * 1000)

while True:
    ret, frame = cap.read()
    if not ret:
        print("摄像头读取失败")
        break

    # 当前时间戳（毫秒）
    now_ms = int(time.time() * 1000)

    # 映射到 90kHz 时间戳（1ms = 90 ticks）
    tmp = now_ms - start_time_ms
    pts = tmp * 90
    print(f"pts: {tmp}, {pts}")

    # 构造 VideoFrame 并设置时间戳
    video_frame = av.VideoFrame.from_ndarray(frame, format="bgr24")
    video_frame.pts = pts
    video_frame.time_base = stream.time_base

    for packet in stream.encode(video_frame):
        container.mux(packet)

    # 停止条件：10秒
    if now_ms - start_time_ms > 10_000:
        break

# flush 编码器
for packet in stream.encode():
    container.mux(packet)

container.close()
cap.release()
cv2.destroyAllWindows()
